#include "spectral_3d.h"

static fftw_complex 		*psi, *ft, *ftp, *pot, *dpsi, *nnl, *ft_nnl, *d2psi, *ftp_exact, *rho_0, *rho_h, *rho_1, *phi_0, *phi_1, *psi_0, *psi_h, *psi_1, *temp;
static double                   *rx, *ry, *rz, *kx, *ky, *kz;
static unsigned int 		nx, ny, nz, ig, jg, kg;
static double 			lx, ly, lz, dt;
static char 			logfile[80];
static fftw_plan                ft_for, ft_bak, ftp_for, ftp_bak, ft_inplace, ift_inplace, ft_psi2nl, ift_ftnnl2nnl, ift_for, ft_temp, ift_temp, ft_0th, ift_0th, ft_ht1, ift_ht1;
static double 			Npc;
static unsigned int		heat;



void ZS_step(){
  
  step_rho_1();
  step1_psi();
  phi_step();
  step_rho_2();
  step2_psi();
  
  for (ig = 0; ig < nx*ny*nz; ig++){
    phi_0[ig] = phi_1[ig];
    psi_0[ig] = psi_1[ig];
    rho_0[ig] = rho_1[ig];
  }
}
void phi_step(){
  unsigned int N = nx*ny*nz;
  
  z_lapl();
  for (ig = 0; ig < N; ig++){
    phi_1[ig] = phi_0[ig] + dt*temp[ig];  
  }
}
void z_lapl(){
    unsigned int N = nx*ny*nz;	

    for (ig = 0; ig < N; ig++){
      temp[ig] = rho_h[ig] - conj(psi_h[ig])*psi_h[ig];
    }
    fftw_execute(ft_temp);
    for (ig = 0; ig < nx; ig++) 
      for (jg = 0; jg < ny; jg++) 
	for (kg = 0; kg < nz; kg++) temp[kg+nz*(jg+ny*ig)] = -1.*(pow(kx[ig],2)+pow(ky[jg],2)+pow(kz[kg],2))*temp[kg+nz*(jg+ny*ig)]/N;
	
    fftw_execute(ift_temp);
}
void step_rho_1(){
  for (ig = 0; ig < nx*ny*nz; ig++) rho_h[ig] = rho_0[ig] + 0.5*dt*phi_0[ig];  
}
void step_rho_2(){
  for (ig = 0; ig < nx*ny*nz; ig++) rho_1[ig] = rho_0[ig] + 0.5*dt*(phi_0[ig] + phi_1[ig]);
}
void step1_psi (){
  
  for (ig = 0; ig < nx*ny*nz; ig++) psi_h[ig] = psi_0[ig];
  z_linear_step1(0.25); 
  z_hstep(0.5);  
  z_linear_step2(0.25);
  
}
void step2_psi (){
  
  z_linear_step1(0.5);
  z_hstep2(1);
  z_linear_step2(0.5);
  for (ig = 0; ig < nx*ny*nz; ig++) psi_1[ig] = psi_h[ig];

}
void z_hstep (fftw_complex a) {
  // evolve one step of laser-plasma interaction

  for (ig = 0; ig < nx*ny*nz; ig ++) {
    psi_h[ig] = psi_h[ig]*cexp(0.5I*a*dt*(rho_0[ig]+rho_h[ig]));
  }
}
void z_hstep2 (fftw_complex a) {
  // evolve one step of laser-plasma interaction

  for (ig = 0; ig < nx*ny*nz; ig ++) {
    psi_h[ig] = psi_h[ig]*cexp(0.5I*a*dt*(rho_0[ig]+rho_1[ig]));
  }
}
void z_linear_step1 (fftw_complex a) {
  // evolves half step of linear Schrodinger equation with no potential exactly
	unsigned int N = nx*ny*nz;	

	fftw_execute(ft_0th);
	for (ig = 0; ig < nx; ig++) 
	   for (jg = 0; jg < ny; jg++) 
	      for (kg = 0; kg < nz; kg++) ft[kg+nz*(jg+ny*ig)] = exp(-1.I*a*dt*(pow(kx[ig],2)+pow(ky[jg],2)+pow(kz[kg],2)))*ft[kg+nz*(jg+ny*ig)]/N;
	      			
	fftw_execute(ift_0th);
}
void z_linear_step2 (fftw_complex a) {
  // evolves half step of linear Schrodinger equation with no potential exactly
	unsigned int N = nx*ny*nz;	

	fftw_execute(ft_0th);
	for (ig = 0; ig < nx; ig++) 
	   for (jg = 0; jg < ny; jg++) 
	      for (kg = 0; kg < nz; kg++) ft[kg+nz*(jg+ny*ig)] = exp(-1.I*a*dt*(pow(kx[ig],2)+pow(ky[jg],2)+pow(kz[kg],2)))*ft[kg+nz*(jg+ny*ig)]/N;
	      			
	fftw_execute(ift_0th);
}
double hamiltonian (){
    double H = 0;
    unsigned int N = nx*ny*nz;
    double dx, dy, dz;

    fftw_execute(ft_for);
    for (ig = 0; ig < nx; ig++) {
      for (jg = 0; jg < ny; jg++) {
	for (kg = 0; kg < nz; kg++) {
	  ft[kg+nz*(jg+ny*ig)] = -1.*(pow(kx[ig],2)+pow(ky[jg],2)+pow(kz[kg],2))*ft[kg+nz*(jg+ny*ig)]/N;  
	}		
      }	
    }
    fftw_execute(ft_bak);
    
    dx = lx/nx; dy = ly/ny; dz = lz/nz;	
    for (ig = 0; ig < nx*ny*nz; ig++ ){
      H = H - dx*dy*dz*( conj(psi[ig])*d2psi[ig]  +  0.5*pow(cabs(psi[ig]),4));
    }
    return H;
} 
double particles() {
  // compute number of particles on grid
  double N = 0.;
  double dx, dy, dz;
  
  dx = lx/nx; dy = ly/ny; dz = lz/nz;
  for (ig = 0; ig < nx*ny*nz; ig++) {
    N = N + dx*dy*dz*pow(cabs(psi[ig]),2);
  }
  return N;
}
double z_particles() {
  // compute number of particles on grid
  double N = 0.;
  double dx, dy, dz;
  
  dx = lx/nx; dy = ly/ny; dz = lz/nz;
  for (ig = 0; ig < nx*ny*nz; ig++) {
    N = N + dx*dy*dz*pow(cabs(psi_0[ig]),2);
  }
  return N;
}
void nnl_step (fftw_complex a) {
  // evolve one step of nonlocal interaction
  convolute();
  if (heat == 1){
    Npc = sqrt(particles()/5.);
  } else Npc = 1.;
  for (ig = 0; ig < nx*ny*nz; ig ++) {
    psi[ig] = psi[ig]*cexp(1.I*a*dt*nnl[ig])/Npc;
  }

}
void nl_step (fftw_complex a) {
  // evolve one step of local |psi|^2 psi interaction

  for (ig = 0; ig < nx*ny*nz; ig ++) {
    psi[ig] = psi[ig]*cexp(1.I*a*dt*pow(cabs(psi[ig]),2));
  }
}
void potential_step() {
// evolve one step update due to potential U(r) 	
  for (ig = 0; ig < nx*ny*nz; ig++){
    psi[ig] = psi[ig]*cexp(1.I*dt*pot[ig]);
  }
}
void convolute() {
  // compute nonlocal potential in Fourier space
  unsigned int N = nx*ny*nz;		
	
  for (ig = 0; ig < N; ig++) {
    dpsi[ig] = pow(cabs(psi[ig]),2);
  }
  fftw_execute(ft_psi2nl);  // here
  for (ig = 0; ig < nx; ig++)
    for (jg = 0; jg < ny; jg++)
      for(kg = 0; kg < nz; kg++){
	ft_nnl[kg+nz*(jg+ny*ig)] = (lx*ly*lz/N)*exp(-0.5*I*(lx*kx[ig]*(1+1./nx) + ly*ky[jg]*(1+1./ny) + lz*kz[kg]*(1+1./nz)))*ftp[kg+nz*(jg+ny*ig)]*ft_nnl[kg+nz*(jg+ny*ig)]/N;  // should be ones
	// ftp  -- numerical, ftp_exact  -- analytic
      }
  fftw_execute(ift_ftnnl2nnl);
}
void linear_step (fftw_complex a) {
  // evolves half step of linear Schrodinger equation with no potential exactly
	unsigned int N = nx*ny*nz;	

	fftw_execute(ft_for);
	for (ig = 0; ig < nx; ig++) {
	   for (jg = 0; jg < ny; jg++) {
	      for (kg = 0; kg < nz; kg++) {
		ft[kg+nz*(jg+ny*ig)] = exp(-1.I*a*dt*(pow(kx[ig],2)+pow(ky[jg],2)+pow(kz[kg],2)))*ft[kg+nz*(jg+ny*ig)]/N;
	      }		
	   }	
	}
	fftw_execute(ift_for);
}	
void init_ft () {
	ft_psi2nl  		= fftw_plan_dft_3d(nx,ny,nz,dpsi,ft_nnl ,FFTW_FORWARD,FFTW_MEASURE);
	ift_ftnnl2nnl 		= fftw_plan_dft_3d(nx,ny,nz,ft_nnl,nnl,FFTW_BACKWARD,FFTW_MEASURE);
	
	ft_for  		= fftw_plan_dft_3d(nx,ny,nz,psi,ft ,FFTW_FORWARD,FFTW_MEASURE);
	ift_for  		= fftw_plan_dft_3d(nx,ny,nz,ft,psi ,FFTW_BACKWARD,FFTW_MEASURE);

	ftp_for 		= fftw_plan_dft_3d(nx,ny,nz,pot,ftp,FFTW_FORWARD,FFTW_MEASURE);
	ft_bak  		= fftw_plan_dft_3d(nx,ny,nz,ft,d2psi,FFTW_BACKWARD,FFTW_MEASURE);
	ftp_bak 		= fftw_plan_dft_3d(nx,ny,nz,ftp,pot,FFTW_BACKWARD,FFTW_MEASURE);
	
	//fftw_plan_with_nthreads(omp_get_num_threads()); // do fft on 2 threads
	ft_inplace 		= fftw_plan_dft_3d(nx,ny,nz,psi,psi,FFTW_BACKWARD,FFTW_MEASURE);
	ift_inplace 		= fftw_plan_dft_3d(nx,ny,nz,psi,psi,FFTW_FORWARD,FFTW_MEASURE);

	ift_ftnnl2nnl 		= fftw_plan_dft_3d(nx,ny,nz,ft_nnl,nnl,FFTW_BACKWARD,FFTW_MEASURE);
	fftw_execute( ftp_for ); // Compute Fourier of potential
	ft_temp			= fftw_plan_dft_3d(nx,ny,nz,temp,temp,FFTW_FORWARD,FFTW_MEASURE);
	ift_temp		= fftw_plan_dft_3d(nx,ny,nz,temp,temp,FFTW_BACKWARD,FFTW_MEASURE);
	
	ft_0th			= fftw_plan_dft_3d(nx,ny,nz,psi_h,ft,FFTW_FORWARD,FFTW_MEASURE);
	ift_0th                 = fftw_plan_dft_3d(nx,ny,nz,ft,psi_h,FFTW_BACKWARD,FFTW_MEASURE);
	
	ft_ht1			= fftw_plan_dft_3d(nx,ny,nz,psi_h,ft,FFTW_FORWARD,FFTW_MEASURE);
	ift_ht1                 = fftw_plan_dft_3d(nx,ny,nz,ft,psi_1,FFTW_BACKWARD,FFTW_MEASURE);
}
void dump_grid_binary (char *file) {
  
	FILE 			*fh;
	
	fh = fopen(file, "wb");

	fwrite(psi,nx*ny*nz*sizeof(fftw_complex),1,fh);
	fclose(fh);
}
void read_binary(char *file){
  
	FILE 			*fh;
	
	fh = fopen(file,"rb");
	if (fh == NULL) {
	  printf("No such data file \n");
	  exit(1);
	} else {
	  if (fread(psi,nx*ny*nz*sizeof(fftw_complex),1,fh) != 1) {
	    printf("Corrupt datafile, probably size mismatch\n");
	    exit(1);
	  } else {
	    printf("Successfully read restart data\n");
	  }
	}
	
	fclose(fh);
}
void dump_fourier_binary (char *file){
	FILE 			*fh;
	fh = fopen(file, "wb");
	
	fwrite(ft,nx*ny*nz*sizeof(fftw_complex),1,fh);
	fclose(fh);
      
}
void dump_fourier (char *file) {

	unsigned int 		i, j, k;	
	char 			dump_file[80];
	FILE			*fh;

	strcpy(dump_file, file);
	strcat(dump_file, ".dat");
	fh = fopen(dump_file,"w");
	fprintf(fh,"  #k_x\t\t  k_y\t\t  k_z\t\t  |ft(psi)|\n\n");
	for ( i = 0; i < nx/2; i++ ) {  // only print the center slice x = 0
	fprintf(fh,"\n\n");
	  for ( j = 0; j < ny/2; j++ ) { 
	  fprintf(fh,"\n");
            for ( k = 0; k < nz/2; k++ ) {
	      fprintf(fh,"%e\t%e\t%e\t%e\n",kx[i],ky[j],kz[k],cabs(ft[k+nz*(j+ny*i)]));
	    }
	  }
	}
	fclose(fh);
}
void dump_grid (char *file) {

	unsigned int 		i, j, k;	
	char 			dump_file[80];
	FILE			*fh;

	strcpy(dump_file, file);
	strcat(dump_file, ".dat");
	fh = fopen(dump_file,"w");
	fprintf(fh,"  #r_x\t\t  r_y\t\t  r_z\t\t  |psi|\t\t  U\n\n");
	for ( i = 0; i < nx; i++ ) {  // only print the center slice x = 0
	fprintf(fh,"\n\n");
	  for ( j = 0; j < ny; j++ ) { 
	  fprintf(fh,"\n");
            for ( k = 0; k < nz; k++ ) {
	      fprintf(fh,"%e\t%e\t%e\t%e\n",rx[i],ry[j],rz[k],cabs(psi[k+nz*(j+ny*i)]));
	    }
	  }
	}
	fclose(fh);
}
void z_dump_grid (char *file) {

	unsigned int 		i, j, k;	
	char 			dump_file[80];
	FILE			*fh;

	strcpy(dump_file, file);
	strcat(dump_file, ".dat");
	fh = fopen(dump_file,"w");
	fprintf(fh,"  #r_x\t\t  r_y\t\t  r_z\t\t  |psi|\t\t  U\n\n");
	for ( i = 0; i < nx; i++ ) {  // only print the center slice x = 0
	fprintf(fh,"\n\n");
	  for ( j = 0; j < ny; j++ ) { 
	  fprintf(fh,"\n");
            for ( k = 0; k < nz; k++ ) {
	      fprintf(fh,"%e\t%e\t%e\t%e\t%e\n",rx[i],ry[j],rz[k],cabs(psi_0[k+nz*(j+ny*i)]),creal(rho_0[k+nz*(j+ny*i)]));
	    }
	  }
	}
	fclose(fh);
}
void init_grid(struct input *in) {

	unsigned int 		i, j, k;
	double 			dx, dy, dz;
	double 			cx, cy, cz, r2, r4;

	cx = 2*M_PI/lx; cy = 2*M_PI/ly; cz = 2*M_PI/lz; 
	
	strcpy(logfile, in->run);
	dx = lx/nx; dy = ly/ny; dz = lz/nz;
	for ( i = 0; i < nx; i++ ) {
	  rx[i] = -0.5*lx + (i+0.5)*dx ;
	  kx[i] = cx*i;
	  if (i > nx/2) kx[i] = kx[i] - cx*nx;
        }
	for ( j = 0; j < ny; j++ ) {
	  ry[j] = -0.5*ly + (j+0.5)*dy;
	  ky[j] = cy*j;
	  if (j > ny/2) ky[j] = ky[j] - cy*ny;
	}
	for ( i = 0; i < nx; i++ ) 
	  for ( j = 0; j < ny; j++ ) 
            for ( k = 0; k < nz; k++ ) {
	      rz[k] = -0.5*lz + (k+0.5)*dz;   
	      r2    = pow(rx[i],2)+pow(ry[j],2)+pow(rz[k],2);
	      r4    = pow(rx[i],4)+pow(ry[j],4)+pow(rz[k],4);
	
	      rho_0[k+nz*(j+ny*i)] = 1./(lx*ly*lz);
	      phi_0[k+nz*(j+ny*i)] = 0;
	      if ((in->ic_type) == 0){
	        psi[k+nz*(j+ny*i)] = (in->A)*exp( -0.5*(in->alpha)*r2 );
	        psi_0[k+nz*(j+ny*i)] = psi[k+nz*(j+ny*i)];
	      } else if ((in->ic_type) == 1)  {
        	 r2    = pow(rx[i]-0.25*lx,2)+pow(ry[j]-0.25*ly,2)+pow(rz[k]-0.25*lz,2);
		 psi[k+nz*(j+ny*i)] = (in->A)*exp( -0.5*(in->alpha)*r2 );
        	 r2    = pow(rx[i]+0.25*lx,2)+pow(ry[j]+0.25*ly,2)+pow(rz[k]+0.25*lz,2);
		 psi[k+nz*(j+ny*i)] = psi[k+nz*(j+ny*i)] + 0.*(in->A)*exp( -0.5*(in->alpha)*r2 );

	      } else if ((in->ic_type) == 2)  {
		 psi[k+nz*(j+ny*i)] = -(in->A)*exp(1.I*atan2(rz[k],ry[j]))*exp(-0.5*(in->alpha)*pow(sqrt(r2) - 3, 2));
	      }
	      pot[k+nz*(j+ny*i)] = 1./pow(r2+(in->delta_2),0.5*(in->b));    // might want to shift that for comparing with rad_sym code
	      if (in->dell_giov  == 1 ){
		 if (r2 != 0) pot[k+nz*(j+ny*i)] = pot[k+nz*(j+ny*i)]*(7/3 + r4/(r2*r2));
		 else pot[k+nz*(j+ny*i)] = pot[k+nz*(j+ny*i)]*(7/3);
	      }			
	      kz[k] = cz*k;
	      if (k > nz/2) kz[k] = kz[k] - cz*nz;
	      ftp_exact[k+nz*(j+ny*i)] = 4*M_PI/(kx[i]*kx[i] + ky[j]*ky[j] + kz[k]*kz[k] + in->delta_2);   // exact fourier transform for Coulomb potential
	    }
}
void load_consts (struct input *in) {

	unsigned int 		i = 0;   
	char     		line[80], value[80];
	FILE* 			fh = fopen(in->run,"r");
	double 			dx, dy, dz;	
	
	if (fh == NULL){
	  printf("Input file missing\n");
	  exit(1);
	}
	
	while(fgets(line, 80, fh) != NULL){
	    i++;
    	    sscanf(line, "%s", value);
	    if ( i == 1) {
		in->nx = atoi(value); 
	    } else if ( i == 2) {
		in->ny = atoi(value); 
	    } else if ( i == 3) {
		in->nz = atoi(value); 
	    } else if ( i == 4) {
		in->cfl = atof(value);
	    } else if ( i == 5) {
		in->lx  = atof(value);
	    } else if ( i == 6) {
		in->ly  = atof(value);
	    } else if ( i == 7) {
		in->lz  = atof(value);
	    } else if ( i == 8) {
		in->alpha  = atof(value);
	    } else if ( i == 9) {
		in->A  = atof(value);
	    } else if ( i == 10) {
		in->delta_2 = atof(value);
	    } else if ( i == 11) {
		in->Tmax = atof(value);
	    } else if ( i == 12) {
		in->nthreads = atoi(value);
	    } else if ( i == 13) {
		in->b = atof(value);
	    } else if ( i == 14) {
		in->heat = atoi(value);
	    } else if ( i == 15) {
		in->io_num = atoi(value);
	    } else if ( i == 16) {
		in->zakh = atoi(value);
	    } else if ( i == 17) {
		in->dell_giov = atoi(value);
	    } else if ( i == 18) {
		in->ic_type = atoi(value);
	    }
	}	
	
	fclose(fh);
	printf("Running:  %s\n", in->run); heat =       in->heat;
	nx  = in->nx;   lx    = in->lx; dx 	=	lx/nx; 
	ny  = in->ny;	ly    = in->ly;	dy 	=	ly/ny;
	nz  = in->nz;	lz    = in->lz; dz 	= 	lz/nz;
	if ((dx <= dy)&&(dx <= dz))
	 {
		dt = (in->cfl)*dx*dx;
	 } else if ((dy <= dx)&&(dy <= dz))
	 {
		dt = (in->cfl)*dy*dy;
	 } else if ((dz <= dx)&&(dz <= dx)) 
	 {
		dt = (in->cfl)*dz*dz;
	 } 
	in->dt = dt;
  	printf("Read:\ncfl=%f\nlx=%f\tnx=%d\nly=%f\tny=%d\nlz=%f\tnz=%d\n",in->cfl,lx,nx,ly,ny,lz,nz);
  	printf("A=%f\talpha=%f\ndelta^2=%f\nU(r) = r^(-%f)\nTime step set to %f\nFinal time time is set to %f\n",in->A,in->alpha,in->delta_2,in->b,in->dt,in->Tmax);
  	printf("Use # threads (not working properly yet)\t\t%d\nWrite out every %d iterations\nSolve Zakharov System yes/no\t\t\t\t %d\nUse Dell Giovanazzi non-symmetric potential yes/no \t %d\n",in->nthreads,in->io_num,in->zakh,in->dell_giov);
	printf("Initializing arrays\t\t");
	init_arrays();
	init_grid(in);
	printf("Complete\nInitializing grid\t\t");
	init_ft();
	init_grid(in);
	printf("Complete\nWriting ICs\t\t\t");
	sprintf(line, "%s.psi.0000", in->run);
	dump_grid_binary(line);	
	printf("Complete\n");
	
}
void init_arrays () {		
  
     phi_0   		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex)); 
     phi_1   		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));
     
     rho_0   		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));
     rho_h   		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));
     rho_1   		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));
     
     psi   		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));
     psi_h		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));
     psi_0		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));
     psi_1		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));

     temp		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex)); // stores (rho - |psi|^2)
     
     ft_nnl   		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));
     dpsi  		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));
     d2psi  		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));
     ft	   		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));	
     ftp   		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));	
     pot   		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));	
     nnl   		= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));	
     ftp_exact	 	= (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));	
     rx    		= (double*) fftw_malloc(nx * sizeof(double));	
     ry    		= (double*) fftw_malloc(ny * sizeof(double));	
     rz    		= (double*) fftw_malloc(nz * sizeof(double));	
     kx   		= (double*) fftw_malloc(nx * sizeof(double));	
     ky    		= (double*) fftw_malloc(ny * sizeof(double));	
     kz    		= (double*) fftw_malloc(nz * sizeof(double));						
}
void free_arrays () {			
	fftw_free(psi);
	fftw_free(dpsi);
	fftw_free(nnl);
	fftw_free(d2psi);
	fftw_free(ft);
	fftw_free(ftp);
	fftw_free(pot);
	fftw_free(rx);
	fftw_free(ry);
	fftw_free(rz);
	fftw_free(kx);
	fftw_free(ky);
	fftw_free(kz);
}